MODULE sefmm_density

   USE fmm_global_paras
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: sefmm_init_density,            &
             sefmm_free_density,            &
             sefmm_get_density,             &
             sefmm_get_order,               &
             sefmm_get_spectral_density

   TYPE(sefmm_rho_hierarchy), ALLOCATABLE :: rho_at_level(:)
   TYPE(sefmm_box_rho), ALLOCATABLE :: hash(:)
   INTEGER(INTK), PARAMETER :: nlevel = 1

   REAL(REALK), SAVE :: used_memory  ! MB
   REAL(REALK), SAVE :: sefmm_rho_threshold
   REAL(REALK), SAVE :: accumulated_order

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_init_density(scheme,basis,cMO,nocc)

      IMPLICIT NONE
      TYPE(scheme_paras), INTENT(IN) :: scheme
      TYPE(fmm_basis),    INTENT(IN) :: basis
      REAL(REALK),        INTENT(IN) :: cMO(:,:)
      INTEGER(INTK),      INTENT(IN) :: nocc

      REAL(REALK) :: nel, T1, T0

      sefmm_rho_threshold = scheme%sefmm_rho_threshold

      CALL CPU_TIME(T0)
      print *, 'Building rho...'
      CALL sefmm_build_rho(scheme,nocc,basis,cMO)
      CALL CPU_TIME(T1)
      CALL TIMTXT('>>> TIME USED build density', T1-T0, LUPRI)

      print *, 'Transforming rho...'
      CALL sefmm_transform_rho(scheme%grain,nel)
      CALL CPU_TIME(T0)
      CALL TIMTXT('>>> TIME USED transform density', T0-T1, LUPRI)
      print *, "Total # electrons  = ", nel

      CALL sefmm_build_hash

!      WRITE(LUPRI,'(A,F7.2)') " SE-density memory (MB) = ", used_memory

   END SUBROUTINE sefmm_init_density

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_free_density

      IMPLICIT NONE

      TYPE(sefmm_box_rho), POINTER :: item
      INTEGER(INTK) :: l, i

      DO l = 1, nlevel
         item => rho_at_level(l)%head
         DO WHILE (ASSOCIATED(item))
            IF (ASSOCIATED (item%coeffs) ) DEALLOCATE (item%coeffs)
            IF (ASSOCIATED (item%rho) ) DEALLOCATE (item%rho)
            NULLIFY (item%coeffs)
            NULLIFY (item%rho)
            item => item%next
         END DO
         NULLIFY (rho_at_level(l)%head)
      END DO
      DEALLOCATE (rho_at_level)

      DO i = 1, SIZE(hash)
         IF (ASSOCIATED (hash(i)%coeffs) ) DEALLOCATE (hash(i)%coeffs)
         IF (ASSOCIATED (hash(i)%rho) ) DEALLOCATE (hash(i)%rho)
         NULLIFY (hash(i)%coeffs)
         NULLIFY (hash(i)%rho)
      END DO
      DEALLOCATE (hash)

   END SUBROUTINE sefmm_free_density

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_order(box,order)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: box
      INTEGER(INTK), INTENT(OUT) :: order

      REAL(REALK), POINTER :: ptr(:,:,:)

      CALL sefmm_get_rho(1,box,ptr,order)
!      NULLIFY (ptr)

   END SUBROUTINE sefmm_get_order

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_density(box,rho,order)

      IMPLICIT NONE
      INTEGER(INTK),        INTENT(IN)  :: box
      REAL(REALK), POINTER, INTENT(OUT) :: rho(:,:,:)
      INTEGER(INTK),        INTENT(OUT) :: order

      CALL sefmm_get_rho(1,box,rho,order)

   END SUBROUTINE sefmm_get_density

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_spectral_density(box,coeffs,order)

      IMPLICIT NONE
      INTEGER(INTK),        INTENT(IN)  :: box
      REAL(REALK), POINTER, INTENT(OUT) :: coeffs(:,:,:)
      INTEGER(INTK),        INTENT(OUT) :: order

      CALL sefmm_get_rho(2,box,coeffs,order)

   END SUBROUTINE sefmm_get_spectral_density

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_rho(mode,box,ptr,order)

      INTEGER(INTK),        INTENT(IN)  :: mode, box
      REAL(REALK), POINTER, INTENT(OUT) :: ptr(:,:,:)
      INTEGER(INTK),        INTENT(OUT) :: order

      order = hash(box)%order
      IF (order == 0) RETURN

      IF (mode == 1) THEN
         ptr => hash(box)%rho(:,:,:)
      ELSE IF (mode == 2) THEN
         ptr => hash(box)%coeffs(:,:,:)
      ELSE
         CALL fmm_quit('sefmm_get_rho')
      END IF

   END SUBROUTINE sefmm_get_rho

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_build_hash

      IMPLICIT NONE
      TYPE(sefmm_box_rho), POINTER :: item
      INTEGER(INTK) :: i, nbox, filled_boxes

      IF (nlevel /= 1) CALL fmm_quit('sefmm_build_hash')

      nbox = rho_at_level(1)%nbox

      ALLOCATE (hash(nbox))
      print *, 'Memory allocation for box hash (MB) =', nbox*8/1e6
      DO i = 1, nbox
         hash(i)%box = 0
         hash(i)%order = 0
         NULLIFY(hash(i)%rho)
         NULLIFY(hash(i)%coeffs)
      END DO

      filled_boxes = 0
      accumulated_order = 0
      item => rho_at_level(1)%head
      DO WHILE (ASSOCIATED(item))

         IF (item%order > 0) THEN
            filled_boxes = filled_boxes +1
            accumulated_order = accumulated_order + item%order
            hash(item%box)%box = item%box
            hash(item%box)%order = item%order
            hash(item%box)%rho => item%rho(:,:,:)
            hash(item%box)%coeffs => item%coeffs(:,:,:)
         END IF

         item => item%next
      END DO

      print *, 'Accumulated order =', accumulated_order
      print *, 'Total rho boxes =', filled_boxes

   END SUBROUTINE sefmm_build_hash

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_build_rho(scheme,nocc,basis,cMO)

      USE fmm_utils,       ONLY: fmm_loop_progress
      USE fmm_shell_pairs, ONLY: fmm_get_shell_pairs
      USE fefmm_box_utils, ONLY: fefmm_get_boxes,       &
                                 fefmm_init_box_utils

      IMPLICIT NONE
      TYPE(scheme_paras), INTENT(IN) :: scheme
      INTEGER(INTK),      INTENT(IN) :: nocc
      TYPE(fmm_basis),    INTENT(IN) :: basis
      REAL(REALK),        INTENT(IN) :: cMO(:,:)

      TYPE(fmm_sh_pairs), POINTER :: sh_pairs(:)
      REAL(REALK)   :: extent(basis%nshells)
      REAL(REALK)   :: grain
      INTEGER(INTK) :: morder
      INTEGER(INTK) :: order, nbox, i, box

      morder = scheme%FEdim
      grain = scheme%grain
      CALL fmm_get_shell_pairs(basis,sh_pairs)
      CALL fefmm_init_box_utils(grain,sh_pairs)
      CALL fefmm_get_boxes(nbox)

      CALL sefmm_get_shell_extents(basis,extent)

      ALLOCATE (rho_at_level(nlevel))
      DO i = 1, nlevel
         NULLIFY (rho_at_level(i)%head)
         rho_at_level(i)%nbox = 0
      END DO

! Only build level 1 for now...
      rho_at_level(1)%nbox = nbox
      DO box = 1, nbox
         CALL fmm_loop_progress(box,nbox)
         IF (scheme%SEdynamic) THEN
            CALL sefmm_get_dynamic_order(morder,grain,box,basis,extent,order)
            IF (order > 15) print *, box, '(',nbox,')', order
         ELSE
            order = morder
         END IF
         IF (order == 0) CYCLE
         CALL sefmm_build_rho_in_box(1,box,order,nocc,grain,basis,extent,cMO)
      END DO

   END SUBROUTINE sefmm_build_rho

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_transform_rho(grain0,nel)

      USE sefmm_cheb_utils, ONLY: cheb_trans_trace

      IMPLICIT NONE
      REAL(REALK),   INTENT(IN)  :: grain0
      REAL(REALK),   INTENT(OUT) :: nel

      TYPE(sefmm_box_rho), POINTER :: item
      INTEGER(INTK) :: l, order
      REAL(REALK)   :: grain

      nel = zero
      grain = grain0

      DO l = 1, nlevel

         item => rho_at_level(l)%head
         DO WHILE (ASSOCIATED(item))
            IF (MAXVAL(item%rho) > sefmm_rho_threshold) THEN
               CALL sefmm_optimize_rho(item)
               order = item%order
               !ALLOCATE (item%coeffs(order,order,order))
               !CALL cheb_get_coeffs(order,item%rho,item%coeffs)
               nel = nel + cheb_trans_trace(order,grain,item%coeffs)
            ELSE
               item%order = 0
               IF (ASSOCIATED (item%rho) ) DEALLOCATE (item%rho)
               NULLIFY (item%rho)
            END IF
            item => item%next
         END DO

         grain = half*grain

      END DO

   END SUBROUTINE sefmm_transform_rho

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_optimize_rho(item)

      USE sefmm_cheb_utils, ONLY: cheb_get_coeffs,         &
                                  cheb_get_function,       &
                                  cheb_truncate_coeffs

      IMPLICIT NONE
      TYPE(sefmm_box_rho), INTENT(INOUT) :: item

      REAL(REALK), POINTER :: ptr(:,:,:)
      REAL(REALK), TARGET :: tmp(128,128,128)

      INTEGER(INTK) :: order, new_order

      order = item%order
      IF (order > 128 ) CALL fmm_quit('sefmm_optimize_rho')
      tmp(1:order,1:order,1:order) = zero

      ptr => tmp(1:order,1:order,1:order)

      CALL cheb_get_coeffs(order,item%rho,ptr)
      CALL cheb_truncate_coeffs(order,ptr,new_order)

      IF (order /= new_order) THEN
         item%order = new_order
         DEALLOCATE (item%rho)
         ALLOCATE (item%rho(new_order,new_order,new_order))
         CALL cheb_get_function(order,new_order,ptr,item%rho)
      END IF

      ALLOCATE (item%coeffs(new_order,new_order,new_order))
      item%coeffs(:,:,:) = ptr(1:new_order,1:new_order,1:new_order)

      IF (new_order > 15) print *, item%box, new_order

   END SUBROUTINE sefmm_optimize_rho

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_build_rho_in_box(l,box,order,nocc,grain,basis,extent,cMO)

      USE fefmm_box_utils, ONLY: fefmm_box_centre,    &
                                 sefmm_AO_in_box

      IMPLICIT NONE
      INTEGER(INTK),   INTENT(IN) :: l, box, order, nocc
      REAL(REALK),     INTENT(IN) :: grain
      TYPE(fmm_basis), INTENT(IN) :: basis
      REAL(REALK),     INTENT(IN) :: extent(basis%nshells)
      REAL(REALK),     INTENT(IN) :: cMO(:,:)

      REAL(REALK)   :: MO(order,order,order,nocc)
      REAL(REALK)   :: rho(order,order,order)
      REAL(REALK)   :: AO(order,order,order)
      REAL(REALK)   :: sh_cntr(3), box_cntr(3)
      REAL(REALK)   :: tmp, thresh
      INTEGER(INTK) :: m, Ish, IAnglA, p,IL1
      LOGICAL :: empty

      MO = zero   
!      thresh = 10d0*SQRT(sefmm_rho_threshold)
      thresh = sefmm_rho_threshold
      box_cntr = fefmm_box_centre(box)

      shells: DO Ish = 1, basis%nshells

         sh_cntr(:) = basis%Centr(:,basis%KAtom(Ish))
         IF (.NOT. sefmm_AO_in_box(grain,extent(Ish),sh_cntr,box_cntr)) CYCLE

         CALL sefmm_get_primitive_rho(order,grain,Ish,box,basis,rho,empty)
         IF (empty) CYCLE shells

         IAnglA = basis%KType(Ish)
         p = basis%KLoc_Car(ISh)

         AOs: DO IL1 = basis%LtuvMin_Car(IAnglA), basis%LtuvMax_Car(IAnglA)
            p = p+1

            CALL sefmm_project_AO(order,grain,basis,Ish,IL1,box,rho,AO)
            tmp = MAXVAL(ABS(AO))

            MOs: DO m = 1, nocc
               IF (tmp * ABS(cMO(p,m)) < thresh) CYCLE
               MO(:,:,:,m) = MO(:,:,:,m) + cMO(p,m) * AO(:,:,:)
            END DO MOs

         END DO AOs

      END DO shells

      CALL sefmm_update_rho(l,order,nocc,box,MO)


   END SUBROUTINE sefmm_build_rho_in_box

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_allocate_rho_box(order,head)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: order
      TYPE(sefmm_box_rho), POINTER :: head
      TYPE(sefmm_box_rho), POINTER :: new_item

      ALLOCATE (new_item)
      NULLIFY(new_item%next)
      NULLIFY (new_item%coeffs)
      ALLOCATE (new_item%rho(order,order,order))

      ! Insert new item at top of linked list (next to head)
      IF (ASSOCIATED(head)) new_item%next => head
      head => new_item

   END SUBROUTINE sefmm_allocate_rho_box

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_update_rho(l,order,nocc,box,MO)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: l, order, nocc, box
      REAL(REALK),   INTENT(IN) :: MO(order,order,order,nocc)

      REAL(REALK), POINTER :: rho(:,:,:)
      REAL(REALK)   :: thresh
      INTEGER(INTK) :: i,j,k, n

!      thresh = 10d0*SQRT(sefmm_rho_threshold)
      thresh = sefmm_rho_threshold

      CALL sefmm_allocate_rho_box(order,rho_at_level(l)%head)

      rho_at_level(l)%head%box = box
      rho_at_level(l)%head%order = order

      rho => rho_at_level(l)%head%rho(:,:,:)
      rho(:,:,:) = zero

      DO n = 1, nocc

         IF ( MAXVAL(ABS(MO(:,:,:,n))) < thresh ) CYCLE

         DO k = 1, order
            DO j = 1, order
               DO i = 1, order

                  rho(i,j,k) = rho(i,j,k) + MO(i,j,k,n) * MO(i,j,k,n)

               END DO
            END DO
         END DO

      END DO

      ! Doubly occupied MOs
      rho(:,:,:) = two*rho(:,:,:)

   END SUBROUTINE sefmm_update_rho

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_dynamic_order(max_order,grain,box,basis,extent,order)

      IMPLICIT NONE
      REAL(REALK),     INTENT(IN)  :: grain
      INTEGER(INTK),   INTENT(IN)  :: max_order, box
      TYPE(fmm_basis), INTENT(IN)  :: basis
      REAL(REALK),     INTENT(IN)  :: extent(basis%nshells)
      INTEGER(INTK),   INTENT(OUT) :: order

      INTEGER(INTK) :: trial

      trial = 2
      DO WHILE (trial <= max_order)
         CALL sefmm_test_dynamic_order(trial,grain,box,basis,extent,order)
         IF (trial > order) EXIT
         trial = trial*2
      END DO

   END SUBROUTINE sefmm_get_dynamic_order

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_test_dynamic_order(order,grain,box,basis,extent,order2)

      USE fefmm_box_utils,  ONLY: fefmm_box_centre,          &
                                  sefmm_AO_in_box
      USE sefmm_cheb_utils, ONLY: cheb_get_nodes,            &
                                  cheb_get_1D_coeffs,        &
                                  cheb_1D_truncate_coeffs

      IMPLICIT NONE
      INTEGER(INTK),   INTENT(IN)  :: order, box
      REAL(REALK),     INTENT(IN)  :: grain
      TYPE(fmm_basis), INTENT(IN)  :: basis
      REAL(REALK),     INTENT(IN)  :: extent(basis%nshells)
      INTEGER(INTK),   INTENT(OUT) :: order2

      REAL(REALK)   :: expn(order,3)
      REAL(REALK)   :: nodes(order)
      REAL(REALK)   :: coeffs(order)
      REAL(REALK)   :: box_cntr(3), sh_cntr(3)
      REAL(REALK)   :: lo(3), hi(3)
      REAL(REALK)   :: bma(3), bpa(3)
      REAL(REALK)   :: gpt, X
      INTEGER(INTK) :: i,j,k, p,ip1,ip2, Ish, order3

      order2 = 0

      box_cntr = fefmm_box_centre(box)
      CALL cheb_get_nodes(order,nodes)

      lo(:) = box_cntr(:) - half*grain
      hi(:) = box_cntr(:) + half*grain
      bma(:) = half*(hi(:)-lo(:))
      bpa(:) = half*(hi(:)+lo(:))

      shells: DO Ish = 1, basis%nshells

         sh_cntr(:) = basis%Centr(:,basis%KAtom(Ish))
         IF (.NOT. sefmm_AO_in_box(grain,extent(Ish),sh_cntr,box_cntr)) CYCLE

         ip1 = basis%KStart(Ish)
         ip2 = ip1 + basis%KontG(Ish) - 1
   
         primitives: DO p = ip1, ip2
   
            DO i = 1, 3
               DO j = 1, order
                  gpt = nodes(j)*bma(i)+bpa(i)
                  X = gpt - sh_cntr(i)
                  expn(j,i) = EXP(-basis%Expnt(p) * X*X)
               END DO
               IF (MAXVAL(expn(:,i)) < sefmm_rho_threshold) CYCLE primitives
            END DO
   
            DO i = 1, 3
               CALL cheb_get_1D_coeffs(order,expn(:,i),coeffs)
               CALL cheb_1D_truncate_coeffs(order,coeffs,order3)
               order2 = MAX(order2,order3)
               IF (order2 >= order) RETURN
            END DO

         END DO primitives

      END DO shells
   
   END SUBROUTINE sefmm_test_dynamic_order

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_primitive_rho(order,grain,Ish,box,basis,rho,empty)

      USE fefmm_box_utils,  ONLY: fefmm_box_centre
      USE sefmm_cheb_utils, ONLY: cheb_get_nodes

      IMPLICIT NONE
      INTEGER(INTK),   INTENT(IN)  :: order, Ish, box
      REAL(REALK),     INTENT(IN)  :: grain
      TYPE(fmm_basis), INTENT(IN)  :: basis
      REAL(REALK),     INTENT(OUT) :: rho(order,order,order)
      LOGICAL,         INTENT(OUT) :: empty

      REAL(REALK)   :: tmp(order,order,order)
      REAL(REALK)   :: expn(order,3)
      REAL(REALK)   :: nodes(order)
      REAL(REALK)   :: box_cntr(3), sh_cntr(3)
      REAL(REALK)   :: lo(3), hi(3)
      REAL(REALK)   :: bma(3), bpa(3)
      REAL(REALK)   :: gpt, X
      INTEGER(INTK) :: i,j,k, p,ip1,ip2

      empty = .TRUE.
      sh_cntr(:) = basis%Centr(:,basis%KAtom(Ish))

      box_cntr = fefmm_box_centre(box)
      CALL cheb_get_nodes(order,nodes)

      lo(:) = box_cntr(:) - half*grain
      hi(:) = box_cntr(:) + half*grain
      bma(:) = half*(hi(:)-lo(:))
      bpa(:) = half*(hi(:)+lo(:))

      ip1 = basis%KStart(Ish)
      ip2 = ip1 + basis%KontG(Ish) - 1

      rho(:,:,:) = zero

      primitives: DO p = ip1, ip2

         DO i = 1, 3
            DO j = 1, order
               gpt = nodes(j)*bma(i)+bpa(i)
               X = gpt - sh_cntr(i)
               expn(j,i) = EXP(-basis%Expnt(p) * X*X)
            END DO
            IF (MAXVAL(expn(:,i)) < sefmm_rho_threshold) CYCLE primitives
         END DO

         DO k = 1, order
            DO j = 1, order
               DO i = 1, order
                  tmp(i,j,k) = expn(i,1)*expn(j,2)*expn(k,3)
               END DO
            END DO
         END DO

         IF (MAXVAL(tmp) < sefmm_rho_threshold) CYCLE primitives
         tmp = basis%CCoef(p) * tmp
         IF (MAXVAL(tmp) < sefmm_rho_threshold) CYCLE primitives

         rho = rho + tmp
         empty = .FALSE.

      END DO primitives

   END SUBROUTINE sefmm_get_primitive_rho

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_project_AO(order,grain,basis,Ish,IL1,box,rho,AO)

      USE fefmm_box_utils,  ONLY: fefmm_box_centre
      USE sefmm_cheb_utils, ONLY: cheb_get_nodes

      IMPLICIT NONE
      INTEGER(INTK),   INTENT(IN)  :: order
      REAL(REALK),     INTENT(IN)  :: grain
      TYPE(fmm_basis), INTENT(IN)  :: basis
      INTEGER(INTK),   INTENT(IN)  :: Ish, IL1, box
      REAL(REALK),     INTENT(IN)  :: rho(order,order,order)
      REAL(REALK),     INTENT(OUT) :: AO(order,order,order)

      REAL(REALK)   :: facXYZ(order,3)
      REAL(REALK)   :: nodes(order)
      REAL(REALK)   :: box_cntr(3), sh_cntr(3)
      REAL(REALK)   :: lo(3),hi(3), bma(3),bpa(3)
      REAL(REALK)   :: facYZ, gpt, X
      INTEGER(INTK) :: tuv(3)
      INTEGER(INTK) :: i,j,k, ij, n

      sh_cntr(:) = basis%Centr(:,basis%KAtom(Ish))

      CALL cheb_get_nodes(order,nodes)
      box_cntr = fefmm_box_centre(box)

      lo(:) = box_cntr(:) - half*grain
      hi(:) = box_cntr(:) + half*grain
      bma(:) = half*(hi(:)-lo(:))
      bpa(:) = half*(hi(:)+lo(:))

      tuv(1) = basis%Lt(IL1)
      tuv(2) = basis%Lu(IL1)
      tuv(3) = basis%Lv(IL1)

      DO i = 1, 3
         DO n = 1, order
            gpt = nodes(n)*bma(i) + bpa(i)
            X = gpt - sh_cntr(i)
            facXYZ(n,i) = X**tuv(i)
         END DO
      END DO

      DO k = 1, order
         DO j = 1, order
            facYZ = facXYZ(j,2) * facXYZ(k,3)
            DO i = 1, order
               AO(i,j,k) = facXYZ(i,1) * facYZ * rho(i,j,k)
            END DO
         END DO
      END DO

   END SUBROUTINE sefmm_project_AO

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_shell_extents(basis,extent)

      IMPLICIT NONE
      TYPE(fmm_basis), INTENT(IN)  :: basis
      REAL(REALK),     INTENT(OUT) :: extent(:)

      INTEGER(INTK) :: Ish, ip1,ip2,p
      REAL(REALK)   :: tmp, thresh

      thresh = -LOG10(sefmm_rho_threshold)

      DO Ish = 1, basis%nshells
         ip1 = basis%KStart(Ish)
         ip2 = ip1 + basis%KontG(Ish) - 1
         tmp = zero
         DO p = ip1, ip2
            tmp = MAX(tmp,SQRT(LOG_10*thresh/basis%Expnt(p)))
         END DO
         extent(Ish) = tmp 
      END DO

   END SUBROUTINE sefmm_get_shell_extents

!-------------------------------------------------------------------------------

END MODULE sefmm_density

